import gradio as gr
import torch
import json
import random
import os
from PIL import Image
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

# Constants
MODEL_PATH = "./output/tari-product-image"
PROCESSED_IMAGES_DIR = "processed_images/img"
DATA_JSON_PATH = "processed_images/data.json"

# Load model and processor
print("Loading model...")
model = Qwen2_5_VLForConditionalGeneration.from_pretrained(MODEL_PATH, device_map="auto")
processor = AutoProcessor.from_pretrained(MODEL_PATH)
print("Model loaded successfully")

# Load product database
with open(DATA_JSON_PATH, 'r', encoding='utf-8') as f:
    product_database = json.load(f)

def create_product_category_index():
    """Create an index mapping categories to product IDs"""
    category_index = {}
    for product_id, product_info in product_database.items():
        category = product_info.get('label', '').strip()
        if category:
            if category not in category_index:
                category_index[category] = []
            category_index[category].append(product_id)
    return category_index

category_index = create_product_category_index()

def get_similar_products(category, num_products=5):
    """Get random similar products from the same category"""
    print(f"Searching for category: {category}")
    if category not in category_index:
        print(f"Category {category} not found in index")
        return []
    
    product_ids = category_index[category]
    selected_ids = random.sample(product_ids, min(num_products, len(product_ids)))
    
    similar_products = []
    for product_id in selected_ids:
        product_info = product_database[product_id]
        image_path = os.path.join(PROCESSED_IMAGES_DIR, os.path.basename(product_info['url']))
        if os.path.exists(image_path):
            try:
                image = Image.open(image_path)
                similar_products.append({
                    'image': image,
                    'title': product_info.get('title', 'No title'),
                    'label': product_info.get('label', 'No label')
                })
            except Exception as e:
                print(f"Error loading image {image_path}: {e}")
                continue
    
    return similar_products

def process_text(text_input):
    """Process text input"""
    if not text_input or text_input.strip() == "":
        return "Please provide a text description.", None, "No products found."
    
    processed_conversation = [{
        "role": "user",
        "content": [
            {"type": "text", "text": f"'{text_input}'这个商品属于什么类别？请只回答类别名称。"}
        ]
    }]

    inputs = processor.apply_chat_template(
        processed_conversation,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt"
    ).to(model.device)

    output_ids = model.generate(**inputs, max_new_tokens=128)
    generated_ids = [output_ids[i][len(inputs.input_ids[i]):] for i in range(len(output_ids))]
    category = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
    
    similar_products = get_similar_products(category.strip())
    
    # 修改这里：为每个图片创建更简洁的标签
    output_images = [(product['image'], f"{product['title']}\n{product['label']}") 
                    for product in similar_products]
    
    # 创建详细信息文本
    details = "\n\n".join([
        f"商品 {i+1}:\n标题: {product['title']}\n类别: {product['label']}"
        for i, product in enumerate(similar_products)
    ])
    
    return category, output_images, details

def process_image(image_input):
    """Process image input"""
    if image_input is None:
        return "Please provide an image.", None, "No products found."
    
    processed_conversation = [{
        "role": "user",
        "content": [
            {"type": "image", "image": image_input},
            {"type": "text", "text": "这是什么类别的商品？请只回答类别名称。"}
        ]
    }]

    inputs = processor.apply_chat_template(
        processed_conversation,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt"
    ).to(model.device)

    output_ids = model.generate(**inputs, max_new_tokens=128)
    generated_ids = [output_ids[i][len(inputs.input_ids[i]):] for i in range(len(output_ids))]
    category = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
    
    similar_products = get_similar_products(category.strip())
    
    # 修改这里：为每个图片创建更简洁的标签
    output_images = [(product['image'], f"{product['title']}\n{product['label']}") 
                    for product in similar_products]
    
    # 创建详细信息文本
    details = "\n\n".join([
        f"商品 {i+1}:\n标题: {product['title']}\n类别: {product['label']}"
        for i, product in enumerate(similar_products)
    ])
    
    return category, output_images, details

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# 商品分类和推荐系统")
    
    with gr.Tabs():
        # Text Search Tab
        with gr.Tab("文本搜索"):
            with gr.Row():
                with gr.Column():
                    text_input = gr.Textbox(label="商品描述", placeholder="请输入商品描述...")
                    text_submit_btn = gr.Button("搜索")
                with gr.Column():
                    text_category_output = gr.Textbox(label="预测类别")
            
            with gr.Row():
                text_gallery = gr.Gallery(
                    label="相似商品",
                    show_label=True,
                    elem_id="text_gallery",
                    columns=5,  # 固定5列
                    rows=1,
                    height=300,  # 设置固定高度
                    object_fit="contain",
                    show_download_button=False,  # 隐藏下载按钮
                )
            
            text_labels_output = gr.Textbox(
                label="商品详细信息",
                visible=True,
                lines=10
            )
            
            text_submit_btn.click(
                fn=process_text,
                inputs=[text_input],
                outputs=[text_category_output, text_gallery, text_labels_output]
            )

        # Image Search Tab
        with gr.Tab("图片搜索"):
            with gr.Row():
                with gr.Column():
                    image_input = gr.Image(label="商品图片", type="pil")
                    image_submit_btn = gr.Button("搜索")
                with gr.Column():
                    image_category_output = gr.Textbox(label="预测类别")
            
            with gr.Row():
                image_gallery = gr.Gallery(
                    label="相似商品",
                    show_label=True,
                    elem_id="image_gallery",
                    columns=5,  # 固定5列
                    rows=1,
                    height=300,  # 设置固定高度
                    object_fit="contain",
                    show_download_button=False,  # 隐藏下载按钮
                )
            
            image_labels_output = gr.Textbox(
                label="商品详细信息",
                visible=True,
                lines=10
            )
            
            image_submit_btn.click(
                fn=process_image,
                inputs=[image_input],
                outputs=[image_category_output, image_gallery, image_labels_output]
            )

# Launch the demo
if __name__ == "__main__":
    demo.launch(share=True)